try:
    import angr
except:
    print("angr is not installed. Please install it by running 'pip install angr'.")
    pass
import json
import os
import logging
import signal

from features.nodes import CFG_Node, CDG_Node, CG_Node, DDG_Node, Edge

from dataclasses import asdict


class TimeoutException(Exception):
    pass


def timeout_handler(signum, frame):
    raise TimeoutException()


# 设置日志文件路径
log_dir = './log/angr'  # 请替换为你的日志保存路径
os.makedirs(log_dir, exist_ok=True)  # 如果日志文件夹不存在则创建
log_file_path = os.path.join(log_dir, 'analysis_errors.log')


# 设置日志记录
logging.basicConfig(
    filename=log_file_path,  # 错误日志文件路径
    level=logging.ERROR,
    format='%(asctime)s - %(levelname)s - %(message)s',
    force=True
)


def get_address_space(binary_path):
    project = angr.Project(binary_path, auto_load_libs=False)
    main_obj = project.loader.main_object
    
    # 主二进制文件的最小和最大地址
    min_addr = main_obj.min_addr
    max_addr = main_obj.max_addr

    return min_addr, max_addr


def get_disasm_info(block):
    """从基本块中提取反汇编信息"""
    disasm_info = []
    for insn in block.capstone.insns:
        disasm_info.append({
            "address": hex(insn.address),
            "mnemonic": insn.mnemonic,
            "op_str": insn.op_str
        })
    return disasm_info


def analyses_cfg(binary_path, min_addr, max_addr):
    # # 设置超时
    # signal.signal(signal.SIGALRM, timeout_handler)
    # signal.alarm(10)  # 设置超时时间为5秒

    try:
        # 初始化项目
        project = angr.Project(binary_path, auto_load_libs=False)

        # 定义初始变量
        function_addrs = set()  # 记录所有已分析的函数地址
        unique_addrs = set()  # 用于去重节点地址
        newly_discovered_addrs = set([project.entry])  # 初始化从入口点开始 project.entry 0x4020d0
        call_depth = 5  # 初始分析的深度

        cfg = None  # 存储最完整的 CFG
        temp_cfg = None

        # 持续分析直到没有新的函数发现
        while newly_discovered_addrs:
            # print(f"开始新一轮分析，分析深度: {call_depth}")
            current_discovered_addrs = set()  # 当前轮次新发现的函数地址

            # 对每个未分析的函数地址进行分析
            for addr in newly_discovered_addrs:
                if min_addr <= addr <= max_addr and addr not in unique_addrs:    
                    # print(f"分析函数地址：{hex(addr)}")
                    temp_cfg = project.analyses.CFGEmulated(
                        context_sensitivity_level=2,
                        keep_state=True,
                        call_depth=call_depth,
                        starts=[addr],  # 从特定函数地址开始分析
                        state_add_options=angr.sim_options.refs
                    )
                    unique_addrs.add(addr)

                    # 如果是第一次执行分析，则保存此 CFG
                    if cfg is None:
                        cfg = temp_cfg
                    else:
                        # 将新 CFG 的函数、块、调用信息合并到现有的 CFG 中
                        # 手动合并函数，因为 FunctionManager 没有 update 方法
                        for func_addr, func in temp_cfg.kb.functions.items():
                            if func_addr not in cfg.kb.functions:
                                cfg.kb.functions[func_addr] = func

                    # 记录新发现的函数地址
                    for func in temp_cfg.kb.functions.values():
                        if func.addr not in function_addrs:
                            current_discovered_addrs.add(func.addr)

            # 更新已分析的函数地址
            function_addrs.update(newly_discovered_addrs)

            # 更新要分析的新函数地址
            newly_discovered_addrs = current_discovered_addrs - function_addrs
            # print(f"发现新的函数数量：{len(newly_discovered_addrs)}")

            # 如果没有新函数发现，可以选择加深分析深度
            if not newly_discovered_addrs:
                break
    except TimeoutException:
        #print("########")
        logging.error(f"Timeout while processing {binary_path}")
        return None, None
    # finally:
    #     signal.alarm(0)  # 取消超时

    return project, cfg


def generate_cfg(project, cfg, min_addr, max_addr, output_path=None):    
    nodes = []
    edges = []

    unique_node_addrs = set()  # 用于去重节点地址
    unique_edges = set()  # 用于去重边 (src_addr, dst_addr)

    """获取节点信息"""
    for node in cfg.graph.nodes():
        if hasattr(node, 'addr') and node.addr is not None:
            if min_addr <= node.addr <= max_addr and node.addr not in unique_node_addrs:
                try:
                    block = project.factory.block(node.addr)
                    func = cfg.kb.functions.get(node.addr)

                    node_info = {
                        'id': node.addr,
                        'addr': hex(node.addr),
                        'function_name': func.name if func else "unknown_func",
                        'function_addr': hex(func.addr) if func else "0x0",
                        'block_addr': hex(block.addr),
                        'disasm': get_disasm_info(block)
                    }
                    
                    node_obj=CFG_Node(**node_info)
                    
                    nodes.append(node_obj)
                    unique_node_addrs.add(node.addr)  # 记录已处理的节点地址

                except Exception as e:
                    print(f"无法处理节点 {hex(node.addr)}: {e}")

    """获取边信息"""
    for src_node in cfg.graph.nodes():
        for dst_node in cfg.graph.successors(src_node):
            if min_addr <= src_node.addr <= max_addr and min_addr <= dst_node.addr <= max_addr:
                edge_data = cfg.graph.get_edge_data(src_node, dst_node)
                edge_key = (src_node.addr, dst_node.addr)  # 用 (src_addr, dst_addr) 作为边的唯一标识
                #print(dir(edge_data.get))
                if edge_key not in unique_edges:
                    edge_info = {
                        'src_id': src_node.addr,
                        'dst_id': dst_node.addr,
                        'type': edge_data.get('jumpkind', 'Contol_Flow')  #边默认属性为 Contol_Flow
                        # 根据需要添加更多属性
                    }
                    
                    edge_obj=Edge(**edge_info)
                    
                    edges.append(edge_obj)   
                    unique_edges.add(edge_key)  # 记录已处理的边  

    # 对节点按照地址进行排序
    sorted_nodes = sorted(nodes, key=lambda x: x.id)            

    if output_path is not None:
        with open(output_path, 'w') as f:
            json.dump({'nodes': [asdict(node) for node in sorted_nodes],'edges': [asdict(edge) for edge in edges]}, f, indent=2)
        print(f"CFG 已成功导出到 {output_path}")

    return sorted_nodes, edges

    
def generate_cdg(project, cfg, min_addr, max_addr, output_path=None):
    cdg = project.analyses.CDG(cfg)
    
    nodes = []
    edges = []

    unique_node_addrs = set()  # 用于去重节点地址
    unique_edges = set()  # 用于去重边 (src_addr, dst_addr)

    """获取节点信息"""
    for node in cdg.graph.nodes():
        if hasattr(node, 'addr') and node.addr is not None:
            if min_addr <= node.addr <= max_addr and node.addr not in unique_node_addrs:
                try:
                    block = project.factory.block(node.addr)
                    func = cfg.kb.functions.get(node.addr)

                    node_info = {
                        'id': node.addr,
                        'addr': hex(node.addr),
                        'function_name': func.name if func else "unknown_func",
                        'function_addr': hex(func.addr) if func else "0x0",
                        'block_addr': hex(block.addr),
                        'disasm': get_disasm_info(block)
                    }
                    
                    node_obj=CDG_Node(**node_info)
                    
                    nodes.append(node_obj)
                    unique_node_addrs.add(node.addr)  # 记录已处理的节点地址

                except Exception as e:
                    print(f"无法处理节点 {hex(node.addr)}: {e}")

    """获取边信息"""
    for src_node, dst_node, data in cdg.graph.edges(data=True):
        if min_addr <= src_node.addr <= max_addr and min_addr <= dst_node.addr <= max_addr:
            edge_key = (src_node.addr, dst_node.addr)  # 用 (src_addr, dst_addr) 作为边的唯一标识
            if edge_key not in unique_edges:
                edge_info = {
                    'src_id': src_node.addr,
                    'dst_id': dst_node.addr,
                    'type': data.get('type', 'Contol_Dependence'),  #边默认属性为 Contol_Dependence
                    # 如果需要更多属性，可以在此添加
                }
                
                edge_obj=Edge(**edge_info)
                
                edges.append(edge_obj)
                unique_edges.add(edge_key)  # 记录已处理的边 
    
    # 对节点按照地址进行排序
    sorted_nodes = sorted(nodes, key=lambda x: x.id) 

    if output_path is not None:
        with open(output_path, 'w') as f:
            json.dump({'nodes': [asdict(node) for node in sorted_nodes],'edges': [asdict(edge) for edge in edges]}, f, indent=2)
        print(f"CDG 已成功导出到 {output_path}")
    return sorted_nodes, edges
   

def generate_ddg(project, cfg, min_addr, max_addr, output_path=None):
    ddg = project.analyses.DDG(cfg)
    
    nodes = []
    edges = []

    unique_node_addrs = set()  # 用于去重节点地址
    unique_edges = set()  # 用于去重边 (src_addr, dst_addr)

    """获取节点信息"""
    for node in ddg.graph.nodes():
        if hasattr(node, 'ins_addr') and node.ins_addr is not None:
            if min_addr <= node.ins_addr <= max_addr and node.ins_addr not in unique_node_addrs:
                try:
                    block = project.factory.block(node.ins_addr)
                    func = cfg.kb.functions.get(node.ins_addr)

                    node_info = {
                        'id': node.ins_addr,
                        'addr': hex(node.ins_addr),
                        'function_name': func.name if func else "unknown_func",
                        'function_addr': hex(func.addr) if func else "0x0",
                        'block_addr': hex(block.addr),
                        'disasm': get_disasm_info(block)
                    }
                    
                    node_obj=DDG_Node(**node_info)
                    
                    nodes.append(node_obj)
                    unique_node_addrs.add(node.ins_addr)  # 记录已处理的节点地址

                except Exception as e:
                    print(f"无法处理节点 {hex(node.ins_addr)}: {e}")

    """获取边信息"""
    for src_node, dst_node, data in ddg.graph.edges(data=True):
        if min_addr <= src_node.ins_addr <= max_addr and min_addr <= dst_node.ins_addr <= max_addr:
            edge_key = (src_node.ins_addr, dst_node.ins_addr)  # 用 (src_addr, dst_addr) 作为边的唯一标识
            if edge_key not in unique_edges:
                edge_info = {
                    'src_id': src_node.ins_addr,
                    'dst_id': dst_node.ins_addr,
                    'type': data.get('type', 'Data_Dependence'),  #边默认属性为  Data_Dependence
                    # 如果需要更多属性，可以在此添加
                }
                
                edge_obj=Edge(**edge_info)
                
                edges.append(edge_obj)
                unique_edges.add(edge_key)  # 记录已处理的边 
                    
    # 对节点按照地址进行排序
    sorted_nodes = sorted(nodes, key=lambda x: x.id) 

    if output_path is not None:
        with open(output_path, 'w') as f:
            json.dump({'nodes': [asdict(node) for node in sorted_nodes],'edges': [asdict(edge) for edge in edges]}, f, indent=2)
        print(f"DDG 已成功导出到 {output_path}")

    return sorted_nodes, edges
    

def generate_cg(project, cfg, min_addr, max_addr, output_path=None):
    cg = cfg.kb.callgraph
    
    nodes = []
    edges = []
    func_addr_to_id = {}
    node_id = 0

    unique_nodes = set()  # 用于去重节点
    unique_edges = set()  # 用于去重边 (src_addr, dst_addr)

    """获取节点信息"""
    for func_addr in cg.nodes():
        if min_addr <= func_addr <= max_addr and func_addr not in unique_nodes:  # 检查是否已经处理过该节点
            func = project.kb.functions.get(func_addr)
            func_name = func.name if func else f"sub_{func_addr:x}"
            # disasm_info = []

            # if func:
            #     for block_addr in func.block_addrs:
            #         try:
            #             block = project.factory.block(block_addr)
            #             disasm_info = get_disasm_info(block)

            #         except Exception as e:
            #             print(f"无法获取块信息: {e} (地址: {hex(block_addr)})")
            block = project.factory.block(func_addr)
            disasm_info = get_disasm_info(block)
            
            node_info = {
                'id': func_addr,
                'addr': hex(func_addr),
                'function_name': func_name,
                'function_addr': hex(func_addr),
                'block_addr': hex(func_addr),  # 假设CG函数地址即为基本块地址
                'disasm': disasm_info
            }
            
            node_obj=CG_Node(**node_info)
            
            nodes.append(node_obj)


            func_addr_to_id[func_addr] = node_id
            node_id += 1

            # 将节点地址添加到去重集合中
            unique_nodes.add(func_addr)

    """获取边信息"""
    for src_addr, dst_addr in cg.edges():
        if min_addr <= src_addr <= max_addr and min_addr <= dst_addr <= max_addr:
            edge_key = (src_addr, dst_addr)  # 使用 (src_addr, dst_addr) 作为边的唯一标识
            if edge_key not in unique_edges:
                if src_addr in func_addr_to_id and dst_addr in func_addr_to_id:
                    
                    edge_info = {
                        'src_id': src_addr,
                        'dst_id': dst_addr,
                        'type': "Call_Function"   #边默认属性为 Call_Function
                    }
                    
                    edge_obj=Edge(**edge_info)
                    
                    edges.append(edge_obj)
                    unique_edges.add(edge_key)  # 将边添加到去重集合中

    # 对节点按照地址进行排序
    sorted_nodes = sorted(nodes, key=lambda x: x.id) 

    if output_path is not None:
        with open(output_path, 'w') as f:
            json.dump({'nodes': [asdict(node) for node in sorted_nodes],'edges': [asdict(edge) for edge in edges]}, f, indent=2)
        print(f"CG 已成功导出到 {output_path}")

    return sorted_nodes, edges


def generate_all_graphs(project, cfg, output_dir, min_addr, max_addr, graph_types):
    os.makedirs(output_dir, exist_ok=True)

    # 生成需要的图
    if 'cfg' in graph_types:
        generate_cfg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'cfg_graph.json'))
    if 'cdg' in graph_types:
        generate_cdg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'cdg_graph.json'))
    if 'ddg' in graph_types:
        generate_ddg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'ddg_graph.json'))
    if 'cg' in graph_types:
        generate_cg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'cg_graph.json'))
    if 'all' in graph_types:
        generate_cfg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'cfg_graph.json'))
        # generate_cdg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'cdg_graph.json'))
        # generate_ddg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'ddg_graph.json'))
        generate_cg(project, cfg, min_addr, max_addr, os.path.join(output_dir, 'cg_graph.json'))


def batch_processing(base_input_dir, output_dir, graph_types):
    for root, dirs, files in os.walk(base_input_dir):
        for filename in files:
            binary_path = os.path.join(root, filename)

            try:
                print(f"正在分析 {binary_path}")
                # 进行地址空间分析
                min_addr, max_addr = get_address_space(binary_path)
                # 进行CFG分析
                # project, cfg = analyses_cfg(binary_path, min_addr, max_addr)
                regions = [(min_addr, max_addr)]
                project = angr.Project(binary_path, auto_load_libs=False)  # 加载二进制文件
                cfg = project.analyses.CFGFast(regions = regions)  # 获取 CFG

                if cfg is None:  # 如果cfg为None，跳过处理
                    continue

                # 创建输出文件夹，保留原有层次结构
                relative_path = os.path.relpath(root, base_input_dir)
                if filename.endswith('.out'):
                    file_output_dir = os.path.join(output_dir, relative_path, filename[:-4])  # 去掉.out后缀
                else:
                    file_output_dir = os.path.join(output_dir, relative_path, filename)  
                
                print(file_output_dir)
                
                # 生成所有图
                generate_all_graphs(project, cfg, file_output_dir, min_addr, max_addr, graph_types)

            except Exception as e:
                # 记录错误信息到日志文件
                logging.error(f"Error processing {binary_path}: {e}")
                
    print("所有分析已完成。")


if __name__ == "__main__":
    # 源文件夹路径
    base_input_dir = '/home/kali/Desktop/angr/test/CWE121_Stack_Based_Buffer_Overflow'
    # 目标输出路径
    output_dir = '/home/kali/Desktop/angr/testgraph/CWE121_Stack_Based_Buffer_Overflow'  # 请替换为你的目标路径

    # 选择需要生成的图
    graph_types = ['cfg', 'cdg', 'ddg']  # 可以选择 ['cfg', 'cdg', 'ddg', 'cg']

    # 执行批处理
    batch_processing(base_input_dir, output_dir, graph_types)
